I wonder, in what cases would you want that the model params also use the same dtype? I have read that when using bfloat16, you usually always want that the params are also stored in bfloat16.

With float16 AMP training, is it normal that you would keep the params in float32?

What about TensorFloat32?

So I wonder if this is sth which should also automatically be handled via torch_amp_options or whether that should be a separate option. And I wonder how people usually do it.

I just merge this as an initial version now.

Can you comment on my questions?

Another question: Was it by intention that you allowed that the user does not specify dtype? I also wonder, if you use dtype=None, what exactly is the behavior of autocast?

With float16 AMP training, is it normal that you would keep the params in float32?

When entering an autocast-enabled region, Tensors may be any type. You should not call half() or bfloat16() on your model(s) or inputs when using autocasting.
https://pytorch.org/docs/stable/amp.html#torch.autocast

So yes, you should not do anything to the model explicitly because autocast is handling that.

Another question: Was it by intention that you allowed that the user does not specify dtype? I also wonder, if you use dtype=None, what exactly is the behavior of autocast?

No, this is a mistake, it should be given. I do not know what the behavior is then, but likely not what is intended.

So yes, you should not do anything to the model explicitly because autocast is handling that.

Autocast is automatically casting inputs to certain PyTorch ops. The parameters are not changed, they are just casted automatically for those ops. But this is not really my question. My question is: Wouldn't it make more sense to directly have the parameters in float16?

I renamed torch_amp_options just to torch_amp.

I was even thinking about renaming it to just amp, or maybe also autocast. We also don't have torch_ prefix for other things, and you might want to use this also for other backends. For example, jmp implements automatic mixed precision training for JAX.

So yes, you should not do anything to the model explicitly because autocast is handling that.

Autocast is automatically casting inputs to certain PyTorch ops. The parameters are not changed, they are just casted automatically for those ops. But this is not really my question. My question is: Wouldn't it make more sense to directly have the parameters in float16?

I see no indication why, unless you really want your whole network to run in float16.

My question is: Wouldn't it make more sense to directly have the parameters in float16?

I see no indication why

Because that further reduce the memory requirement? Why would you not want that? What are the downsides?

I'm not saying that everything should be float16. Maybe certain ops must stay in float32. I thought this is the main aspect of autocast/AMP, to cast to float16 wherever it makes sense.

I just don't understand why weights are stored in float32, and then always auto-casted. That also adds some overhead in computation (the casting), and requires more memory. Unless there is maybe some reason. But that is my question, what is the reason for this?

Ah, I was just checking the original paper introducing automatic mixed precision training, and it explains it (Sec 3.1):

In mixed precision training, weights, activations and gradients are stored as FP16. In order to match the accuracy of the FP32 networks, an FP32 master copy of weights is maintained and updated with the weight gradient during the optimizer step. In each iteration an FP16 copy of the master weights is used in the forward and backward pass. ...

While the need for FP32 master weights is not universal, there are two possible reasons why a number of networks require it. One explanation is that updates (weight gradients multiplied by the learning rate) become too small to be represented in FP16 - any value whose magnitude is smaller than $2^{−24}$ becomes zero in FP16. ...

Another explanation is that the ratio of the weight value to the weight update is very large. In this case, even though the weight update is representable in FP16, it could still become zero when addition operation right-shifts it to align the binary point with the weight. ...

Now that the code is longer, we might want to move this into the updater or add an extra module instead of having it plain in the engine.